#! python
#-*- coding: utf-8 -*-


def list_commands():
    return ['list', 'generate', 'chat']


def parse(text: str):
    index = 0
    if text[index] == "{":
        length, key = parse(text[1:])
        index = index + length
        assert(isinstance(key, str))
        while text[index] == " ":
            index = index + 1
        if text[index] == ":":
            length, value = parse(text[index+1:])
            index = index + length
    elif text[0] == "[":
        parse(text)
    # print(f"parse({text[0]})")


def json_decode(in_bytes, encoding="utf-8"):
    import json
    if isinstance(in_bytes, bytes):
        text = in_bytes.decode(encoding)
    else:
        # print(f"type(in_bytes) = {type(in_bytes)}")
        text = in_bytes
    # print(f"type(text) = {type(text)}")
    return parse(text)


def generate(prompt, url=None, headers=None, data=None, **kwargs):
    import requests
    url = url or "http://localhost:11434/api/generate"
    headers = headers or { "Content-Type": "application/json" }
    data = data or { "model": "deepseek-r1:1.5b", "prompt": prompt }
    import json
    data = json.dumps(data)
    response = requests.post(url=url, headers=headers, data=data)
    response.raise_for_status()
    # 输出响应内容
    for chunk in response.iter_lines():
        if chunk:
            decoded_chunk = json.loads(chunk.decode('utf-8'))
            if decoded_chunk['done']:
                print(f'\n', flush=True)
                return {
                        "done_reason": decoded_chunk["done_reason"],
                        "load_duration": decoded_chunk["load_duration"],
                        "total_duration": decoded_chunk["total_duration"], 
                        "prompt_eval_count": decoded_chunk["prompt_eval_count"],
                        "prompt_eval_duration": decoded_chunk["prompt_eval_duration"],
                        "eval_count": decoded_chunk["eval_count"],
                        "eval_duration": decoded_chunk["eval_duration"]
                        }
            else:
                print(f'{decoded_chunk["response"]}', end="", flush=True)
    return response


def main():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('command', help='command to execute',
                        choices=list_commands())
    parser.add_argument('prompt', help='prompt to models', default='test')
    args = parser.parse_args()
    if (args.command == 'list'):
        print(f"available commands are: \n{'\n'.join(list_commands())}")
    if (args.command == 'generate'):
        print(generate(args.prompt))


if __name__ == '__main__':
    main()
